import math
import time
from pathlib import Path
from typing import List

from PIL import Image


def list_files_with_pathlib(folder_path: str or Path, recursive: bool = True) -> List[Path]:
    """
    获取指定文件夹下的所有文件。

    :param folder_path: 文件夹路径（字符串或 Path 对象）
    :param recursive: 是否递归遍历子目录，默认为 True
    :return: 包含所有文件 Path 对象的列表
    """
    folder = Path(folder_path)

    if not folder.exists():
        raise FileNotFoundError(f"路径不存在: {folder}")

    if not folder.is_dir():
        raise NotADirectoryError(f"指定的路径不是文件夹: {folder}")

    # 根据是否递归选择方法
    pattern = '**/*' if recursive else '*'

    # 获取所有匹配的项目，并过滤出文件
    files = [item for item in folder.glob(pattern) if item.is_file()]

    return sorted(files)  # 按名称排序返回，便于查看


if __name__ == '__main__':
    print('开始')
    timestamp = int(time.time())
    image_1 = []
    folder_path = "D:\\1\\"
    files = list_files_with_pathlib(folder_path)
    for file_1 in files:
        image_1.append(str(file_1)[0:-8])
    unique_arr = list(set(image_1))
    print(len(unique_arr))
    for index, aa in enumerate(unique_arr):
        if index % 1000 == 0:
            print(str(index) + "   " + str(len(unique_arr)))
        cols = 0
        rows = 0
        image_paths = []
        # 根据文件名找到对应的图片组
        output_path = aa + ".png"
        output_path = output_path.replace('\\1\\', '\\2\\')
        files = list_files_with_pathlib(folder_path)
        for file_1 in files:
            if str(file_1).startswith(aa):
                image_paths.append(str(file_1))
                temp_num = int(str(file_1)[-8:-7])
                if temp_num > rows:
                    rows = temp_num
        rows = (rows + 1)
        cols = math.ceil(len(image_paths) / rows)
        images = [Image.open(img) for img in image_paths]
        # 找出最大宽高（统一尺寸可选）
        max_width = max(img.width for img in images)
        max_height = max(img.height for img in images)
        total_width = max_width * cols
        total_height = max_height * rows
        merged = Image.new('RGBA', (total_width, total_height))
        for idx in range(len(image_paths)):
            wei = idx // rows
            shou = idx % rows
            shou = int(shou)
            aaa = ""
            if wei > 99:
                aaa = aa + str(shou) + str(wei) + ".png"
            elif wei > 9:
                aaa = aa + str(shou) + "0" + str(wei) + ".png"
            elif wei <= 9:
                aaa = aa + str(shou) + "00" + str(wei) + ".png"
            img = Image.open(aaa)
            x = cols * max_width
            y = rows * max_height
            x_offset = max_width * wei
            y_offset = max_height * shou
            merged.paste(img, (x_offset, y_offset), img)
        dst = Path(output_path)
        dst.parent.mkdir(parents=True, exist_ok=True)
        merged.save(output_path)

    endTime = int(time.time())
    print(str(endTime - timestamp))
    print('结束')
